# everything_search.py
from mcp.server.fastmcp import FastMCP
import sys
import logging
import os
import ctypes
import platform
import time

logger = logging.getLogger('EverythingSearch')

# Fix UTF-8 encoding for Windows console
if sys.platform == 'win32':
    sys.stderr.reconfigure(encoding='utf-8')
    sys.stdout.reconfigure(encoding='utf-8')

# 确定系统架构并加载相应的DLL
bitness = platform.architecture()[0]
if bitness == '64bit':
    dll_path = os.path.join(os.path.dirname(__file__), 'sdk', 'dll', 'Everything64.dll')
else:
    dll_path = os.path.join(os.path.dirname(__file__), 'sdk', 'dll', 'Everything32.dll')

# 加载Everything DLL
try:
    everything_dll = ctypes.WinDLL(dll_path)
    logger.info(f"Successfully loaded {dll_path}")
except Exception as e:
    logger.error(f"Failed to load Everything DLL: {e}")
    sys.exit(1)

# 定义常量
everything_constants = {
    # 错误代码
    'EVERYTHING_OK': 0,
    'EVERYTHING_ERROR_MEMORY': 1,
    'EVERYTHING_ERROR_IPC': 2,
    
    # 排序选项
    'EVERYTHING_SORT_NAME_ASCENDING': 1,
    'EVERYTHING_SORT_NAME_DESCENDING': 2,
    'EVERYTHING_SORT_PATH_ASCENDING': 3,
    'EVERYTHING_SORT_PATH_DESCENDING': 4,
    'EVERYTHING_SORT_SIZE_ASCENDING': 5,
    'EVERYTHING_SORT_SIZE_DESCENDING': 6,
    'EVERYTHING_SORT_DATE_MODIFIED_ASCENDING': 13,
    'EVERYTHING_SORT_DATE_MODIFIED_DESCENDING': 14,
    
    # 请求标志
    'EVERYTHING_REQUEST_FILE_NAME': 0x00000001,
    'EVERYTHING_REQUEST_PATH': 0x00000002,
    'EVERYTHING_REQUEST_FULL_PATH_AND_FILE_NAME': 0x00000004,
    'EVERYTHING_REQUEST_EXTENSION': 0x00000008,
    'EVERYTHING_REQUEST_SIZE': 0x00000010,
    'EVERYTHING_REQUEST_DATE_CREATED': 0x00000020,
    'EVERYTHING_REQUEST_DATE_MODIFIED': 0x00000040,
    'EVERYTHING_REQUEST_DATE_ACCESSED': 0x00000080,
    'EVERYTHING_REQUEST_ATTRIBUTES': 0x00000100
}

# 定义DLL函数
everything_dll.Everything_SetSearchW.argtypes = [ctypes.c_wchar_p]
everything_dll.Everything_SetMatchPath.argtypes = [ctypes.c_bool]
everything_dll.Everything_SetMatchCase.argtypes = [ctypes.c_bool]
everything_dll.Everything_SetMatchWholeWord.argtypes = [ctypes.c_bool]
everything_dll.Everything_SetRegex.argtypes = [ctypes.c_bool]
everything_dll.Everything_SetMax.argtypes = [ctypes.c_ulong]
everything_dll.Everything_SetOffset.argtypes = [ctypes.c_ulong]
everything_dll.Everything_SetSort.argtypes = [ctypes.c_ulong]
everything_dll.Everything_SetRequestFlags.argtypes = [ctypes.c_ulong]
everything_dll.Everything_QueryW.argtypes = [ctypes.c_bool]
everything_dll.Everything_GetLastError.argtypes = []
everything_dll.Everything_GetLastError.restype = ctypes.c_ulong
everything_dll.Everything_GetNumResults.argtypes = []
everything_dll.Everything_GetNumResults.restype = ctypes.c_ulong
everything_dll.Everything_GetNumFileResults.argtypes = []
everything_dll.Everything_GetNumFileResults.restype = ctypes.c_ulong
everything_dll.Everything_GetNumFolderResults.argtypes = []
everything_dll.Everything_GetNumFolderResults.restype = ctypes.c_ulong
everything_dll.Everything_GetTotResults.argtypes = []
everything_dll.Everything_GetTotResults.restype = ctypes.c_ulong
everything_dll.Everything_IsFileResult.argtypes = [ctypes.c_ulong]
everything_dll.Everything_IsFileResult.restype = ctypes.c_bool
everything_dll.Everything_IsFolderResult.argtypes = [ctypes.c_ulong]
everything_dll.Everything_IsFolderResult.restype = ctypes.c_bool
everything_dll.Everything_GetResultFileNameW.argtypes = [ctypes.c_ulong]
everything_dll.Everything_GetResultFileNameW.restype = ctypes.c_wchar_p
everything_dll.Everything_GetResultPathW.argtypes = [ctypes.c_ulong]
everything_dll.Everything_GetResultPathW.restype = ctypes.c_wchar_p
everything_dll.Everything_GetResultFullPathNameW.argtypes = [ctypes.c_ulong, ctypes.c_wchar_p, ctypes.c_ulong]
everything_dll.Everything_GetResultFullPathNameW.restype = ctypes.c_ulong
everything_dll.Everything_Reset.argtypes = []
everything_dll.Everything_CleanUp.argtypes = []
everything_dll.Everything_IsDBLoaded.argtypes = []
everything_dll.Everything_IsDBLoaded.restype = ctypes.c_bool

class EverythingSearch:
    @staticmethod
    def search(query, match_path=False, match_case=False, match_whole_word=False,
               regex=False, max_results=100, sort=1, timeout=10):
        """使用Everything SDK进行文件搜索，带超时控制"""
        try:
            # 重置之前的搜索状态
            everything_dll.Everything_Reset()
            
            # 设置搜索参数 - 简化请求标志，先确保基本功能正常
            everything_dll.Everything_SetSearchW(query)
            everything_dll.Everything_SetMatchPath(match_path)
            everything_dll.Everything_SetMatchCase(match_case)
            everything_dll.Everything_SetMatchWholeWord(match_whole_word)
            everything_dll.Everything_SetRegex(regex)
            everything_dll.Everything_SetMax(max_results)
            everything_dll.Everything_SetSort(sort)
            
            # 使用基本的请求标志
            request_flags = (
                everything_constants['EVERYTHING_REQUEST_FILE_NAME'] |
                everything_constants['EVERYTHING_REQUEST_PATH'] |
                everything_constants['EVERYTHING_REQUEST_FULL_PATH_AND_FILE_NAME']
            )
            everything_dll.Everything_SetRequestFlags(request_flags)
            
            # 先检查Everything服务是否运行
            error = everything_dll.Everything_GetLastError()
            if error != everything_constants['EVERYTHING_OK']:
                if error == everything_constants['EVERYTHING_ERROR_IPC']:
                    return {"success": False, "error": "Everything搜索客户端未运行，请先启动Everything程序"}
                else:
                    return {"success": False, "error": f"初始化错误: Everything API error {error}"}
            
            # 使用阻塞模式查询，这是更安全的方式
            start_time = time.time()
            success = everything_dll.Everything_QueryW(True)
            
            # 检查查询是否成功
            if not success:
                error = everything_dll.Everything_GetLastError()
                error_messages = {
                    1: "内存不足",
                    2: "Everything搜索客户端未运行",
                    7: "无效调用，请检查API使用顺序",
                    8: "无效请求数据"
                }
                error_msg = error_messages.get(error, f"搜索失败: Everything API error {error}")
                return {"success": False, "error": error_msg}
            
            # 检查是否超时
            if time.time() - start_time >= timeout:
                return {"success": False, "error": "搜索在10秒后超时"}
            
            # 获取搜索结果
            results = []
            num_results = everything_dll.Everything_GetNumResults()
            
            # 安全地处理结果
            if num_results > 0:
                for i in range(min(num_results, max_results)):
                    try:
                        # 获取文件名和路径
                        file_name = everything_dll.Everything_GetResultFileNameW(i)
                        path = everything_dll.Everything_GetResultPathW(i)
                        
                        # 获取完整路径
                        buf = ctypes.create_unicode_buffer(260)
                        everything_dll.Everything_GetResultFullPathNameW(i, buf, ctypes.sizeof(buf))
                        full_path = buf.value
                        
                        # 构建结果项
                        result_item = {
                            "full_path": full_path,
                            "file_name": file_name,
                            "path": path,
                            "is_file": everything_dll.Everything_IsFileResult(i),
                            "is_folder": everything_dll.Everything_IsFolderResult(i)
                        }
                        results.append(result_item)
                    except Exception as item_error:
                        logger.warning(f"处理结果项时出错: {item_error}")
                        continue
            
            # 获取统计信息
            stats = {
                "total_results": everything_dll.Everything_GetTotResults(),
                "file_results": everything_dll.Everything_GetNumFileResults(),
                "folder_results": everything_dll.Everything_GetNumFolderResults(),
                "returned_results": len(results)
            }
            
            return {
                "success": True,
                "query": query,
                "stats": stats,
                "results": results
            }
            
        except Exception as e:
            logger.error(f"搜索过程中出错: {e}")
            return {"success": False, "error": f"异常错误: {str(e)}"}
        finally:
            # 清理资源
            everything_dll.Everything_CleanUp()

# 创建MCP服务器
mcp = FastMCP("EverythingSearch")

@mcp.tool()
def search(query: str, match_path: bool = False, match_case: bool = False,
          match_whole_word: bool = False, regex: bool = False,
          max_results: int = 100, sort_by: str = "name") -> dict:
    """
    使用Everything搜索引擎进行快速文件搜索。
    
    参数:
    - query: 搜索查询字符串
    - match_path: 是否匹配完整路径 (默认: False)
    - match_case: 是否大小写敏感 (默认: False)
    - match_whole_word: 是否匹配整个单词 (默认: False)
    - regex: 是否使用正则表达式 (默认: False)
    - max_results: 最大返回结果数 (默认: 100)
    - sort_by: 排序方式，可以是 'name', 'path', 'size', 'date' (默认: 'name')
    
    返回:
    - 包含搜索结果的字典
    """
    # 映射排序参数到SDK常量
    sort_map = {
        "name": everything_constants['EVERYTHING_SORT_NAME_ASCENDING'],
        "path": everything_constants['EVERYTHING_SORT_PATH_ASCENDING'],
        "size": everything_constants['EVERYTHING_SORT_SIZE_ASCENDING'],
        "date": everything_constants['EVERYTHING_SORT_DATE_MODIFIED_ASCENDING']
    }
    
    sort_value = sort_map.get(sort_by.lower(), everything_constants['EVERYTHING_SORT_NAME_ASCENDING'])
    
    # 记录搜索请求
    logger.info(f"Search request: query='{query}', match_path={match_path}, match_case={match_case}, max_results={max_results}")
    
    # 执行搜索，超时设置为10秒
    result = EverythingSearch.search(
        query=query,
        match_path=match_path,
        match_case=match_case,
        match_whole_word=match_whole_word,
        regex=regex,
        max_results=max_results,
        sort=sort_value,
        timeout=10  # 确保超时不超过10秒
    )
    
    # 记录搜索结果
    if result["success"]:
        logger.info(f"Search completed: found {result['stats']['total_results']} results, returned {result['stats']['returned_results']}")
    else:
        logger.warning(f"Search failed: {result.get('error', 'Unknown error')}")
    
    return result

@mcp.tool()
def get_status() -> dict:
    """
    获取Everything服务的状态信息。
    
    返回:
    - 包含服务状态的字典
    """
    try:
        # 检查DLL加载状态
        if 'everything_dll' not in globals():
            return {"success": False, "status": "DLL not loaded"}
        
        # 检查错误状态
        error = everything_dll.Everything_GetLastError()
        
        # 检查数据库加载状态
        db_loaded = everything_dll.Everything_IsDBLoaded()
        
        return {
            "success": True,
            "status": {
                "dll_loaded": True,
                "everything_running": error == everything_constants['EVERYTHING_OK'],
                "db_loaded": db_loaded,
                "bitness": bitness
            }
        }
        
    except Exception as e:
        logger.error(f"Error getting status: {e}")
        return {"success": False, "error": str(e)}

# 启动服务器
if __name__ == "__main__":
    mcp.run(transport="stdio")